#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np
import os

def linspace_custom(start, end, num, output_dtype=np.float32):
    """自定义linspace实现:输入start/end为float16,输出为float32"""
    if num <= 1:
        return np.array([start], dtype=output_dtype)
    start_f32 = start.astype(np.float32)
    end_f32 = end.astype(np.float32)
    step = (end_f32 - start_f32) / (num - 1)
    seq = start_f32 + np.arange(num, dtype=np.float32) * step
    return seq.astype(output_dtype)

def gen_golden_data_linspace():
    np.random.seed()

    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)
    
    # ==================== 固定数据类型配置 ====================
    # input_dtype = np.float16    
    input_dtype = np.float32
    # input_dtype = np.int8
    # input_dtype = np.int16
    # input_dtype = np.int32
    output_dtype = np.float32   
    ns_dtype = np.uint32         
    # =======================基础参数配置======================
    min_n = 1                       # n的最小值
    max_n = 1024                    # n的最大值
    start_range = (-10.0, 10.0)     # start的取值范围
    end_offset_range = (0.1, 20.0)  # end与start的差值范围
    # ========================================================
    starts = np.random.uniform(*start_range, size=1).astype(input_dtype)
    end_offset = np.random.uniform(*end_offset_range, size=1).astype(input_dtype)
    ends = starts + end_offset
    ns = np.random.randint(min_n, max_n + 1, size=1, dtype=ns_dtype)
    golden = linspace_custom(starts[0], ends[0], ns[0], output_dtype=output_dtype)
    
    starts.tofile("./input/starts.bin")
    ends.tofile("./input/ends.bin")
    ns.tofile("./input/ns.bin")
    golden.tofile("./output/golden.bin")
    
if __name__ == "__main__":
    gen_golden_data_linspace()